package cn.com.basic.framework.datahub;

import cn.com.basic.framework.datahub.proxy.ReadProxy;
import cn.hutool.core.thread.ThreadFactoryBuilder;
import com.aliyun.datahub.client.DatahubClient;
import com.aliyun.datahub.client.model.OpenSubscriptionSessionResult;
import com.aliyun.datahub.client.model.TupleRecordData;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * 订阅者，向{@link ReadProxy} 设定消费逻辑
 *
 * @author zhuxj
 */
public class DatahubSubscriber {


    private DatahubClient datahubClient;
    /**
     * 订阅项目空间
     */
    private String projectName;
    /**
     * 订阅会话
     */
    private String topicName;
    /**
     * 订阅号
     */
    private String subId;

    private ExecutorService executorService;
    private List<String> shardIds;
    private OpenSubscriptionSessionResult sessionResult;
    private HashMap<String, ReadProxy> reader = new HashMap<>();

    public DatahubSubscriber(DatahubClient client, String projectName, String topicName, String subId) {
        this.datahubClient = client;
        this.projectName = projectName;
        this.topicName = topicName;
        this.subId = subId;

        this.shardIds = new ArrayList<>();
        for (int i = 0; i < this.datahubClient.getTopic(projectName, topicName).getShardCount(); i++) {
            this.shardIds.add(i + "");
        }

        ThreadFactory namedThreadFactory = new ThreadFactoryBuilder()
                .setNamePrefix("dataHub-subsribe-pool-%d")
                .build();
        this.executorService = new ThreadPoolExecutor(shardIds.size(), shardIds.size(),
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>(),namedThreadFactory);

        this.sessionResult = this.datahubClient.openSubscriptionSession(projectName, topicName, subId, this.shardIds);

        for (String shardId : this.shardIds) {
            this.reader.put(shardId, new ReadProxy(this.datahubClient, this.projectName, this.topicName, this.subId));
        }
    }


    /**
     * 注册消费逻辑
     *
     * @param consumer {@link Consumer<TupleRecordData>}
     */
    public void addConsumer(Consumer<TupleRecordData> consumer) {
        reader.forEach((shardId, reader) -> {
            reader.setShardId(shardId);
            reader.setConsumer(consumer);
            reader.setSessionResult(sessionResult);
        });
    }

    /**
     * 开始读取
     */
    public void run() {
        reader.forEach((shardId, reader) -> executorService.execute(reader::startRead));
    }

}
